import numpy as np
from sklearn import linear_model
from Evaluate import Evaluate
from sklearn.ensemble import RandomForestRegressor

class Esembler:
    def __init__(self):
        pass
    def createLinearEsemble(self,results, ground_truth):
        '''
            results: list of results
        '''
        tmp_eval = []
        self.e = []
        results = np.array(results).swapaxes(0, 3)#tasks, paths, tianshu, models,
        #g : paths, tianshu, tasks
        for i in range(results.shape[0]):
            for j in range(results.shape[1]):
                self.e.append(linear_model.LinearRegression())
                self.e[i * results.shape[1] + j].fit(results[i,j], ground_truth[j, :,i])
                # print results[i,j]
                # print ground_truth[j, :, i]
                tmp_eval.append(self.e[i * results.shape[1] + j].predict(results[i,j]))
                # print tmp_eval[-1]


        #tasks, paths, tianshu
        tmp_eval = np.array(tmp_eval).reshape(results.shape[0], results.shape[1], -1).swapaxes(0, 1).swapaxes(1, 2)
        print tmp_eval.shape
        evaluate = Evaluate()
        MAPE = evaluate.evaluate_test(tmp_eval, ground_truth)
        print 'esemble train MAPE=' + str(MAPE)

    def esemble(self, results):
        tmp_eval = []
        results = np.array(results).swapaxes(0, 3)#tasks, paths, tianshu, models,
        #g : paths, tianshu, tasks
        for i in range(results.shape[0]):
            for j in range(results.shape[1]):
                tmp_eval.append(self.e[i * results.shape[1] + j].predict(results[i,j]))
        tmp_eval = np.array(tmp_eval).reshape(results.shape[0], results.shape[1], -1).swapaxes(0, 1).swapaxes(1, 2)
        return tmp_eval